'''
1. 圈复杂度过高。需要完整优化。
2. 部分代码可以并行化处理。
3. 结构过于凌乱。
4. 冗余的for循环逻辑过多。
'''


from dataclasses import dataclass
import warnings
from collections import Counter

import numpy as np
import pandas as pd
from joblib import load
from rdkit import RDLogger, DataStructs
from rdkit.Chem import AllChem, MACCSkeys

import similaritycalc as sim
import logging


class Arr:
    def __init__(self,reactant,pubchem_id):
        self.reactant = reactant
        self.pubchem_id = pubchem_id

        m = AllChem.MolFromSmiles(sim.canonicalize_smiles(reactant))     
        maccs_keys = MACCSkeys.GenMACCSKeys(m)
        array1 = np.zeros((0,), dtype=np.float64)
        DataStructs.ConvertToNumpyArray(maccs_keys, array1)
        fps = AllChem.GetMorganFingerprintAsBitVect(m, 3, nBits=2048)
        array2 = np.zeros((0,), dtype=np.float64)
        DataStructs.ConvertToNumpyArray(fps, array2)
        array3 = np.append(array1, array2, axis=0)


        self.test_arr_2d = np.reshape(array3, (1, 2215))

     

    def analy(self,df_items,mlp):   
        array = mlp.predict_proba(self.test_arr_2d)
        test_df = df_items.drop(df_items[df_items['PUBCHEM CID'] == self.pubchem_id].index)
        logging.debug("test_df len %d",len(test_df))
        
        result_list = []
        for j in range(len(test_df)):
            try:
                logging.debug("test_df len %d",len(test_df['extract templates'][j]))
                product_list = sim.apply_templates(self.reactant, test_df['extract templates'][j])
                if product_list == "":
                    logging.error("apply_templates error")
                logging.debug("%s--------%s-------%d",self.reactant,test_df['extract templates'][j],len(product_list))
                score_r = sim.calculate_similarity(sim.get_fingerprint(self.reactant, "Morgan2"),
                                                   sim.get_fingerprint(test_df['substrate_smiles_canonical'][j],
                                                                       "Morgan2"),
                                                   "Tanimoto")
                reaction_class_pos = int(test_df['ML class no.'][j]) - 1
                # score_ml = array[reaction_class_pos][0][1]
                score_ml = array[0][reaction_class_pos]
               

                # 加速
                reaction_class = test_df['reaction class no.'][j]
                reaction_id = test_df['Reaction ID'][j]
                for product in product_list:
                    result_dict = self.get_score(test_df['prod_smiles_canonical'][j],product,score_ml,score_r)
                    result_dict['class'] = reaction_class
                    result_dict['reaction_id'] = reaction_id
                    result_list.append(result_dict)
            except Exception as e:
                logging.error("exception")
                logging.error(e)
                pass
            result_list_sorted = sorted(result_list, key=lambda a: a['score'], reverse=True)   

            return self.get_result(result_list_sorted)


    def get_result(self,result_list_sorted):
        result_output_list = []
        result_dict_list = []
        for r in result_list_sorted:
            if sim.canonicalize_smiles(r['product']) not in result_output_list:
                lst = []
                for res in result_dict_list:
                    lst.append(res['class'])
                counter = Counter(lst)
                if counter[r['class']] < 30:  # could be adjusted
                    result_output_list.append(sim.canonicalize_smiles(r['product']))
                    result_dict_list.append(r)
        return result_output_list,result_dict_list

                
    ## 抽出函数，可能需要并行化
    def get_score(self,prod_smiles_canonical,product,score_ml,score_r,rate=0.1):
        result_dict = {}
        score_p = sim.calculate_similarity(sim.get_fingerprint(product, "Morgan2"),
                                                       sim.get_fingerprint(prod_smiles_canonical,
                                                                           "Morgan2"),
                                                       "Tanimoto")
        score = score_r * score_p + score_ml * rate
        result_dict['score'] = score
        result_dict['product'] = product
        return result_dict



class Project:
    def __init__(self):
        # ignore warnings
        warnings.filterwarnings('ignore')
        RDLogger.DisableLog('rdApp.*')

        # write log config
        logging.basicConfig(filename='./example.log', level=logging.DEBUG)

        self.clf = load('clf.joblib')
        self.mlp = load('mlp.joblib')


    
    def read_items(self,df_filename="database_0714.xlsx",data_filename="test set 715.xlsx"):
        self.df_items = pd.read_excel(df_filename)
        self.data_items = pd.read_excel(data_filename)

        logging.debug(self.df_items.head(1))
        logging.debug(self.data_items.head(1))

        self.data_items["product_list"] = None
        self.data_items["product_dict_list"] = None


    def handle_items(self):
        for i in range(len(self.data_items)):
            reactant = self.data_items['substrate_smiles_canonical'][i]
            pubchem_id = self.data_items['PUBCHEM CID'][i]

            arr = Arr(reactant,pubchem_id)
            result_output_list,result_dict_list = arr.analy(self.df_items,self.mlp)
            logging.info(len(result_output_list))
            logging.info(i, "completed", result_output_list[:3])
            self.data_items['product_list'][i] = ';'.join(result_output_list)
            self.data_items['product_dict_list'][i] = result_dict_list


        # get accuracy metrics
        self.data_items["sim_ml_limited_class_precision"] = None
        ## 这不是一个bug？
        self.data_items["sim_ml_limited_class_recall"][i] = None  

    def to_excel(self):
        for i in range(len(self.data_items)):
            try:
                product_list = self.data_items['product sum'][i].split(";")
                sim_list = self.data_items['product_list'][i].split(";")
                self.data_items["sim_ml_limited_class_precision"][i] = sim.get_precison(3, sim_list, product_list)
                self.data_items["sim_ml_limited_class_recall"][i] = sim.get_recall(3, sim_list, product_list)
            except (Exception,):
                self.data_items["sim_ml_limited_class_precision"][i] = 0
                self.data_items["sim_ml_limited_class_recall"][i] = 0

        logging.info(self.data_items["sim_ml_limited_class_precision"].mean(), self.data_items["sim_ml_limited_class_recall"].mean())
        self.data_items.to_excel("test715.xlsx")
         


if __name__ == "__main__":
    pj = Project()
    pj.read_items()
    pj.handle_items()
    pj.to_excel()

